package mr

import (
	"6.824/utils"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"time"
)
import "log"
import "net/rpc"
import "hash/fnv"

//
// Map functions return a slice of KeyValue.
//
type KeyValue struct {
	Key   string
	Value string
}
type ByKey []KeyValue

func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

//
// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
//
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

//
// main/mrworker.go calls this function.
//
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	for {
		reply, err := register()
		if err != nil {
			break
		}
		//log.Print("reply code: ",reply.ReplyCode)
		if reply.ReplyCode == 1 {
			if solveMap(reply, mapf) != nil {
				log.Fatal("fail to solve map")
				returnAns(WorkerRetArgs{TaskID: reply.TaskID, WorkMode: 1, ErrCode: -1})
			}

		} else if reply.ReplyCode == 2 {
			if solveReduce(reply, reducef) != nil {
				log.Fatal("fail to solve reduce")
				returnAns(WorkerRetArgs{TaskID: reply.TaskID, WorkMode: 2, ErrCode: -1})
			}
		} else if reply.ReplyCode == 3 {
			time.Sleep(10)
			continue
		} else {
			break
		}

	}
	// uncomment to send the Example RPC to the coordinator.
	// CallExample()

}

func solveMap(reply *RegisterWorkerReply, mapf func(string, string) []KeyValue) error {
	dic := make(map[int][]KeyValue)
	for file := range utils.ReadFiles(reply.FileNames) {
		mapret := mapf(file.Name, file.Content)
		for _, kv := range mapret {
			k := ihash(kv.Key) % reply.InputCnt
			lis, ok := dic[k]
			if !ok {
				lis = make([]KeyValue, 0)
			}
			dic[k] = append(lis, kv)
		}

	}
	for k, v := range dic {
		f, err := ioutil.TempFile("./", "temp")
		if err != nil {
			return err
		}
		enc := json.NewEncoder(f)
		for _, lv := range v {
			enc.Encode(&lv)
		}
		f.Close()
		os.Rename(f.Name(), fmt.Sprintf("mr-%d-%d", reply.TaskID, k))
	}
	return returnAns(WorkerRetArgs{reply.TaskID, 1, 0})
}
func solveReduce(reply *RegisterWorkerReply, reducef func(string, []string) string) error {
	kva := make([]KeyValue, 0)
	for i := 0; i < reply.InputCnt; i++ {
		fn := fmt.Sprintf("mr-%d-%d", i, reply.TaskID)
		f, err := os.Open(fn)
		if err != nil {
			log.Printf("fail to open %s", fn)
			continue
		}

		dec := json.NewDecoder(f)
		for {
			var kv KeyValue
			if err := dec.Decode(&kv); err != nil {
				break
			}
			kva = append(kva, kv)
		}
	}
	sort.Sort(ByKey(kva))

	f, err := ioutil.TempFile("./", "temp")
	if err != nil {
		return errors.New("fail to create temp file")
	}
	for i := 0; i < len(kva); {
		j := i + 1
		for j < len(kva) && kva[j].Key == kva[i].Key {
			j++
		}
		values := []string{}
		for k := i; k < j; k++ {
			values = append(values, kva[k].Value)
		}
		output := reducef(kva[i].Key, values)

		fmt.Fprintf(f, "%v %v\n", kva[i].Key, output)

		i = j

	}
	f.Close()

	os.Rename(f.Name(), fmt.Sprintf("mr-out-%d", reply.TaskID))
	return returnAns(WorkerRetArgs{reply.TaskID, 2, 0})
}
func register() (*RegisterWorkerReply, error) {
	reply := new(RegisterWorkerReply)
	k := call("Coordinator.Register", new(RegisterWorkerArgs), reply)
	if k {
		return reply, nil
	} else {
		reply.ReplyCode = -1
		return reply, errors.New("Register Failed")
	}
}
func returnAns(args WorkerRetArgs) error {
	ok := call("Coordinator.WorkerReturn", &args, new(WorkerRetReply))
	if !ok {
		return errors.New("return error")
	}
	return nil
}

//
// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
//
func CallExample(i int) {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = i

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	// the "Coordinator.Example" tells the
	// receiving server that we'd like to call
	// the Example() method of struct Coordinator.
	ok := call("Coordinator.Example", &args, &reply)
	if ok {
		// reply.Y should be 100.
		fmt.Printf("reply.Y %v\n", reply.Y)
	} else {
		fmt.Printf("call failed!\n")
	}
}

//
// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	c, err := rpc.DialHTTP("tcp", "101.43.122.187"+":6824")
	//sockname := coordinatorSock()
	//c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
